import os
import glob

# create train.txt and test.txt of kaggle ddd data, containing image path and label\

data_root = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection'
fine_tune_root = '/media/lab/labdisk/home/lab1314/xryu/data/DistractedDriverDetection/fine-tune-caffe/'


def create_source_file():
    f_trian = open(fine_tune_root + '/data_txt/train.txt', mode='w')
    for j in range(8):
        print('Load folder c{}'.format(j))
        train_path = os.path.join(data_root, 'imgs', 'train',
                                  'c' + str(j), '*.jpg')
        files_train = glob.glob(train_path)

        for fl_train in files_train:
            # print fl
            line = fl_train + ' ' + str(j) + '\n'
            f_trian.writelines(line)
    f_trian.close()

    f_test = open(fine_tune_root + '/data_txt/test.txt', mode='w')
    #test_path = os.path.join(data_root, 'imgs', 'test', '*.jpg')
    #files_test = glob.glob(test_path)
    #for fl_test in files_test:
    #    f_test.writelines(fl_test + '\n')
    for j in range(8,10):
        print('Load folder c{}'.format(j))
        test_path = os.path.join(data_root, 'imgs', 'train',
                                  'c' + str(j), '*.jpg')
        files_test = glob.glob(test_path)

        for fl_test in files_test:
            # print fl
            line = fl_test + ' ' + str(j) + '\n'
            f_test.writelines(line)

    f_test.close()


create_source_file()
